from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import status
from django.core.files.storage import default_storage, FileSystemStorage
from django.core.files.base import ContentFile
from .service import MnistService, MnistDropoutService
import os

class MnistPredictView(APIView):
    def post(self, request, *args, **kwargs):
        print("开始执行cnn模型进行识别")
        path = uploadFile(request)
        response = executeService(MnistService(), path)
        return response


## 加入Dropout的模型的视图
class MnistPredictDropoutView(APIView):
    def post(self, request, *args, **kwargs):
        print("开始执行cnn_dropout模型进行识别")
        path = uploadFile(request)
        response = executeService(MnistDropoutService(), path)
        return response

def uploadFile(request):
    # 检查是否有文件上传
    if 'file' not in request.FILES:
        return Response({"error": "No file uploaded"}, status=status.HTTP_400_BAD_REQUEST)

    # 获取上传的文件
    uploaded_file = request.FILES['file']
    file_path = os.path.join('static/image', uploaded_file.name)
    # 保存文件到本地
    path = default_storage.save(file_path, ContentFile(uploaded_file.read()))
    return path

def executeService(service, path):
    try:
        # 调用预测函数
        prediction, confidence = service.predict(path)
        response =  Response({
            "prediction": prediction,
            "confidence": confidence
        }, status=status.HTTP_200_OK)
    except Exception as e:
        response =  Response({"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
    finally:
        # 删除临时文件
        if os.path.exists(path):
            print("文件上传成功,路径为:", path)
            os.remove(path)
    return response